package com.zengde.ai_sandbox.service;

import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.chat.StreamingChatModel;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.chat.response.StreamingChatResponseHandler;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.model.openai.OpenAiStreamingChatModel;
import lombok.Getter;
import lombok.Setter;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

import java.time.Duration;

@Service
public class DoubaoCustomService {

    @Value("${ARK_API_KEY}")
    private String apiKey;

    @Value("${my.doubao.url}")
    private String baseUrl;

    @Getter
    @Setter
    private String modelName;



    /**
     * 手动获取豆包的流式模型
     * @return
     */
    private StreamingChatModel getStreamModel() {
        return OpenAiStreamingChatModel.builder()
                .baseUrl(baseUrl)
                .apiKey(apiKey)
                .modelName(modelName)
                .timeout(Duration.ofSeconds(60))
                .build();
    }

    /**
     * 手动获取豆包的非流式模型
     * @return
     */
    private ChatModel getNonStreamModel() {
        return OpenAiChatModel.builder()
                .baseUrl(baseUrl)
                .apiKey(apiKey)
                .modelName(modelName)
                .maxRetries(3)
                .timeout(Duration.ofSeconds(60))
                .build();
    }

    public String chat(String userMessage) {
        return getNonStreamModel().chat(userMessage);
    }


    public Flux<String> streamChat(String userMessage) {

//        getStreamModel().a

        Flux<String> objectFlux = Flux.create(fluxSink -> {
            getStreamModel().chat(userMessage, new StreamingChatResponseHandler() {
                @Override
                public void onPartialResponse(String s) {
                    fluxSink.next(s);
                }

                @Override
                public void onCompleteResponse(ChatResponse chatResponse) {
                    fluxSink.complete();
                }

                @Override
                public void onError(Throwable throwable) {
                    fluxSink.error(throwable);
                }
            });
        });
        return objectFlux;
    }



}
